package com.boxuegu.websecurity.bbs.filter;

import com.alibaba.fastjson.JSON;
import com.sun.org.apache.regexp.internal.RE;
import com.sun.scenario.effect.impl.sw.java.JSWBrightpassPeer;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * @program: cyberDefense
 * @description:
 * @author: fengjd
 * @create: 2022-07-27 13:58
 **/
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    HttpServletRequest orgRequest = null;

    /**
     * Constructs a request object wrapping the given request.
     *
     * @param request The request to wrap
     * @throws IllegalArgumentException if the request is null
     */
    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
        this.orgRequest = request;
    }

    /**
     * 判断 httpServletRequest 是不是 wrapper类,是的话,强转成wapper,调用getorgRequest方法获取原始request的静态方法
     */
    public static HttpServletRequest getOrgRequest(HttpServletRequest request) {
        if (request instanceof XssHttpServletRequestWrapper) {
            return ((XssHttpServletRequestWrapper) request).getOrgRequest();
        }
        return request;
    }

    public HttpServletRequest getOrgRequest() {
        return orgRequest;
    }


    /* 重写 方法来给参数转义 */

    /**
     * 覆盖getParameter方法，将参数名和参数值都做xss过滤。<br/>
     * 如果需要获得原始的值，则通过super.getParameterValues(name)来获取<br/>
     * getParameterNames,getParameterValues和getParameterMap也可能需要覆盖
     *
     * @param name a <code>String</code> specifying the name of the parameter
     * @return
     */
    @Override
    public String getParameter(String name) {
        String value = super.getParameter(xssEncode(name));
        if (value != null) {
            value = xssEncode(value);
        }
        return value;
    }

    /**
     * 重写 getParameterMap方法
     *
     * @return
     */
    @Override
    public Map<String, String[]> getParameterMap() {
        HashMap<String, String[]> parameterMap = (HashMap<String, String[]>) super.getParameterMap();
        HashMap<String, String[]> clone = (HashMap<String, String[]>) parameterMap.clone();
        for (Iterator<Map.Entry<String, String[]>> iterator = clone.entrySet().iterator(); iterator.hasNext(); ) {
            Map.Entry<String, String[]> entry = (Map.Entry<String, String[]>) iterator.next();
            String[] values = entry.getValue();
            for (int i = 0; i < values.length; i++) {
                if (values[i] instanceof String) {
                    values[i] = xssEncode(values[i]);
                }
            }
            entry.setValue(values);
        }
        return clone;
    }

    /**
     * 重写getParameterValues 方法
     *
     * @param name a <code>String</code> containing the name of the parameter
     *             whose value is requested
     * @return
     */
    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        int count = values.length;
        String[] encodedValues = new String[count];
        for (int i = 0; i < count; i++) {
            encodedValues[i] = xssEncode(values[i]);
        }
        return encodedValues;
    }

    /**
     * 覆盖getheader方法,将参数名和参数值都做xss过滤
     * 如果需要获取原始的值,则通过super.getheader(name)来获取
     * getHeaderNames 也可能需要覆盖
     *
     * @param name
     * @return
     */
    public String getHeader(String name) {
        String value = super.getHeader(xssEncode(name));
        if (value != null) {
            value = xssEncode(value);
        }
        return value;
    }

    /**
     * @return
     */
    @Override
    public String getQueryString() {
        return xssEncode(super.getQueryString());
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        String str = getRequestBody(super.getInputStream());
        Map<String, Object> map = JSON.parseObject(str, Map.class);
        Map<Object, Object> resultMap = new HashMap<>(map.size());
        for (String key : map.keySet()) {
            Object val = map.get(key);
            if (map.get(key) instanceof String) {
                resultMap.put(key,xssEncode(val.toString()));
            }else {
                resultMap.put(key,(val));
            }
        }

         str = JSON.toJSONString(resultMap);
         final ByteArrayInputStream bais= new ByteArrayInputStream(str.getBytes());
         return new ServletInputStream() {
             @Override
             public boolean isFinished() {
                 return false;
             }

             @Override
             public boolean isReady() {
                 return false;
             }

             @Override
             public void setReadListener(ReadListener listener) {

             }

             @Override
             public int read() throws IOException {
                 return bais.read();
             }
         };
    }

    public String xssEncode(String xss) {
        return XSSEncode3.xssEncode(xss);
    }

    private String getRequestBody(InputStream stream) {
        String line = "";
        StringBuilder body = new StringBuilder();
        int counter = 0;

        /* 读取post提交的数据 */
        BufferedReader reader = new BufferedReader(new InputStreamReader(stream, Charset.forName("UTF-8")));

        try {
            while ((line = reader.readLine()) != null) {
                body.append(line);
                counter++;
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        return body.toString();
    }
}
